package com.asiainfo.sqlextractor.functions;

import com.google.common.collect.ImmutableList;
import org.apache.calcite.linq4j.tree.Types;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFamily;
import org.apache.calcite.rex.RexBuilder;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.rex.RexUtil;
import org.apache.calcite.schema.ScalarFunction;
import org.apache.calcite.schema.impl.ScalarFunctionImpl;
import org.apache.calcite.sql.*;
import org.apache.calcite.sql.parser.SqlParserPos;
import org.apache.calcite.sql.type.*;
import org.apache.calcite.sql.validate.SqlUserDefinedFunction;
import org.apache.calcite.sql2rel.*;
import org.apache.calcite.util.Pair;
import org.checkerframework.checker.nullness.qual.Nullable;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;

public class StandardConvertletTableExt implements SqlRexConvertletTable {

    public static final StandardConvertletTableExt INSTANCE = new StandardConvertletTableExt();

    private SqlRexConvertletTable target = StandardConvertletTable.INSTANCE;

    private Method registerOpMethod;


    private StandardConvertletTableExt() {
        init();

        registerOp(createSqlFunction("CALL_FUNC"), createSqlRexConvertlet("CALL_FUNC", "callFunc"));
        registerOp(createSqlFunction("CALL_API"), createSqlRexConvertlet("CALL_API", "callApi"));
        registerOp(createSqlFunction("CALL_SCRIPT"), createSqlRexConvertlet("CALL_SCRIPT", "callScript"));
    }


    private void init() {
        try {
            registerOpMethod = ReflectiveConvertletTable.class.getDeclaredMethod("registerOp",
                    SqlOperator.class, SqlRexConvertlet.class);
            registerOpMethod.setAccessible(true);
        } catch (NoSuchMethodException e) {
            e.printStackTrace();
        }
    }

    public SqlRexConvertlet get(SqlCall call) {
        return target.get(call);
    }

    protected void registerOp(SqlOperator op, SqlRexConvertlet convertlet) {
        try {
            registerOpMethod.invoke(target, op, convertlet);
        } catch (IllegalAccessException e) {
        } catch (InvocationTargetException e) {
        }
    }

    public static SqlFunction createSqlFunction(String name) {
        return new SqlFunction(
                name, SqlKind.OTHER_FUNCTION,
                opBinding -> {
                    RelDataType type = opBinding.getOperandType(opBinding.getOperandCount() - 1);
                    return opBinding.getTypeFactory().createJavaType(Object.class);
                }, null,
                OperandTypes.VARIADIC,
                SqlFunctionCategory.USER_DEFINED_FUNCTION
        );
    }

    private static SqlRexConvertlet createSqlRexConvertlet(String sqlIdentifierName, String methodName) {
        return (ctx, call) -> {
            final RexBuilder rexBuilder = ctx.getRexBuilder();
            final List<RexNode> operands = convertOperands(ctx, call, SqlOperandTypeChecker.Consistency.NONE);
            final RelDataType type = ctx.getValidator().getValidatedNodeType(call);
            final ScalarFunction callFunc = ScalarFunctionImpl.create(
                    Types.lookupMethod(Functions.class, methodName, String.class, Object[].class)
            );
            final SqlIdentifier sqlIdentifier = new SqlIdentifier(ImmutableList.of(sqlIdentifierName), SqlParserPos.ZERO);
            final SqlOperandMetadata operandMetadata = OperandTypes.operandMetadata(ImmutableList.of(SqlTypeFamily.ANY),
                    typeFactory -> ImmutableList.of(type), i -> "arg" + i, i -> false);
            final SqlUserDefinedFunction udf = new SqlUserDefinedFunction(sqlIdentifier, SqlKind.OTHER_FUNCTION,
                    operatorBinding -> type, null, operandMetadata, callFunc);
            return rexBuilder.makeCall(type, udf, operands);
        };
    }


    private static List<RexNode> convertOperands(SqlRexContext cx, SqlCall call,
                                                 SqlOperandTypeChecker.Consistency consistency) {
        return convertOperands(cx, call, call.getOperandList(), consistency);
    }

    private static List<RexNode> convertOperands(SqlRexContext cx,
                                                 SqlCall call, List<SqlNode> nodes,
                                                 SqlOperandTypeChecker.Consistency consistency) {
        final List<RexNode> exprs = new ArrayList<>();
        for (SqlNode node : nodes) {
            exprs.add(cx.convertExpression(node));
        }
        final List<RelDataType> operandTypes =
                cx.getValidator().getValidatedOperandTypes(call);
        if (operandTypes != null) {
            final List<RexNode> oldExprs = new ArrayList<>(exprs);
            exprs.clear();
            Pair.forEach(oldExprs, operandTypes, (expr, type) ->
                    exprs.add(cx.getRexBuilder().ensureType(type, expr, true)));
        }
        if (exprs.size() > 1) {
            final RelDataType type =
                    consistentType(cx, consistency, RexUtil.types(exprs));
            if (type != null) {
                final List<RexNode> oldExprs = new ArrayList<>(exprs);
                exprs.clear();
                for (RexNode expr : oldExprs) {
                    exprs.add(cx.getRexBuilder().ensureType(type, expr, true));
                }
            }
        }
        return exprs;
    }


    private static @Nullable RelDataType consistentType(SqlRexContext cx,
                                                        SqlOperandTypeChecker.Consistency consistency, List<RelDataType> types) {
        switch (consistency) {
            case COMPARE:
                if (SqlTypeUtil.areSameFamily(types)) {
                    // All arguments are of same family. No need for explicit casts.
                    return null;
                }
                final List<RelDataType> nonCharacterTypes = new ArrayList<>();
                for (RelDataType type : types) {
                    if (type.getFamily() != SqlTypeFamily.CHARACTER) {
                        nonCharacterTypes.add(type);
                    }
                }
                if (!nonCharacterTypes.isEmpty()) {
                    final int typeCount = types.size();
                    types = nonCharacterTypes;
                    if (nonCharacterTypes.size() < typeCount) {
                        final RelDataTypeFamily family =
                                nonCharacterTypes.get(0).getFamily();
                        if (family instanceof SqlTypeFamily) {
                            // The character arguments might be larger than the numeric
                            // argument. Give ourselves some headroom.
                            switch ((SqlTypeFamily) family) {
                                case INTEGER:
                                case NUMERIC:
                                    nonCharacterTypes.add(
                                            cx.getTypeFactory().createSqlType(SqlTypeName.BIGINT));
                                    break;
                                default:
                                    break;
                            }
                        }
                    }
                }
                // fall through
            case LEAST_RESTRICTIVE:
                return cx.getTypeFactory().leastRestrictive(types);
            default:
                return null;
        }
    }

}

